fstats_mcmc.f90 Source File


Source Code

! SPDX-FileCopyrightText: 2022-2026 Jason Christopherson
! SPDX-License-Identifier: MIT
!
! MIT License
!
! Copyright (c) 2022-2026 Jason Christopherson
!
! Permission is hereby granted, free of charge, to any person obtaining a copy
! of this software and associated documentation files (the "Software"), to deal
! in the Software without restriction, including without limitation the rights
! to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
! copies of the Software, and to permit persons to whom the Software is
! furnished to do so, subject to the following conditions:
!
! The above copyright notice and this permission notice shall be included in all
! copies or substantial portions of the Software.
!
! THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
! IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
! FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
! AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
! LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
! OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
! SOFTWARE.

module fstats_mcmc
    use iso_fortran_env
    use ieee_arithmetic
    use fstats_distributions
    use fstats_normal_distribution, only : normal_distribution
    use fstats_log_normal_distribution, only : log_normal_distribution
    use fstats_sampling
    use fstats_errors
    use fstats_descriptive_statistics
    use fstats_types
    use fstats_regression
    use fstats_sampling
    use collections
    !$ use omp_lib, only : omp_in_parallel
    implicit none
    private
    public :: chain_builder
    public :: mcmc_sampler
    public :: mcmc_target
    public :: evaluate_model
    public :: mcmc_proposal
    public :: sample_chains
    
    type, abstract :: mcmc_target
        !! Defines a model of the target distribution(s).
        type(list), private :: m_items
            !! The list of parameters.
        real(real64), private, allocatable, dimension(:) :: m_y
            !! A workspace array for containing the model values.
        real(real64), public :: data_noise = 1.0d0
            !! A parameter representing the noise in the data.
        integer(int32), public :: likelihood_parallel_threshold = 10000
            !! Minimum observation count for a threaded likelihood reduction.
            !! Adjust for the workload; nested worker teams are not created.
    contains
        procedure(evaluate_model), deferred, public :: model
        procedure, public :: get_parameter_count => mt_get_param_count
        procedure, public :: add_parameter => mt_add_param
        procedure, public :: get_parameter => mt_get_param
        procedure, public :: likelihood => mt_likelihood
        procedure, public :: likelihood_log_variance => mt_log_likelihood
            !! Gaussian log likelihood accepting log variance directly.
        procedure, public :: evaluate_log_variance_prior => mt_eval_log_var_prior
            !! Prior density in log-variance coordinates, including the Jacobian.
        procedure, public :: log_posterior => mt_log_posterior
            !! Posterior density in parameter/log-variance coordinates.
        procedure, public :: evaluate_variance_prior => mt_eval_var_prior
        procedure, public :: sample_variance_prior => mt_sample_var_prior
        procedure, public :: evaluate_prior => mt_eval_prior
    end type

    interface
        subroutine evaluate_model(this, xdata, xc, y)
            !! Evaluates the model at the supplied values.
            use iso_fortran_env, only : real64
            import mcmc_target
            class(mcmc_target), intent(in) :: this
                !! The mcmc_target object.
            real(real64), intent(in), dimension(:) :: xdata
                !! An M-element array containing the values at which to evaluate
                !! the model.
            real(real64), intent(in), dimension(:) :: xc
                !! An N-element array containing the model parameters.
            real(real64), intent(out), dimension(:) :: y
                !! An M-element array where the resulting model values wil
                !! be written.
        end subroutine
    end interface

! ------------------------------------------------------------------------------
    type :: mcmc_proposal
        !! Defines a type responsible for generating a proposal state for a
        !! Monte-Carlo, Markov-Chain sampler.
        logical, private :: m_recenter = .true.
            !! Allow recentering?
        real(real64), private :: m_scale = 0.1d0
            !! Global proposal scale multiplier (applied to per-parameter stddev)
        real(real64), private, allocatable, dimension(:) :: m_param_scale
            !! Optional per-parameter scale multipliers. If not allocated, m_scale is used.
    contains
        procedure, public :: generate_sample => mp_gen
        procedure, public :: get_recenter => mp_get_recenter
        procedure, public :: set_recenter => mp_set_recenter
        procedure, public :: get_scale => mp_get_scale
        procedure, public :: set_scale => mp_set_scale
        procedure, public :: get_param_scale => mp_get_param_scale
        procedure, public :: set_param_scale => mp_set_param_scale
        procedure, public :: set_param_scales => mp_set_param_scales
    end type

! ------------------------------------------------------------------------------
    type chain_builder
        !! A type allowing for the construction of chain of values.
        integer(int32), private :: initial_iteration_estimate = 10000
            !! An initial estimate at the number of allowed iterations.
        integer(int32), private :: m_bufferSize = 0
            !! The actual number of states (# of used rows) in the buffer.
        real(real64), private, allocatable, dimension(:,:) :: m_buffer
            !! The buffer where each new state is stored as a row in the matrix.
        integer(int32), private :: m_numVars = 0
            !! The number of state variables.
    contains
        procedure, public :: get_state_variable_count => cb_get_nvars
        procedure, public :: get_chain_length => cb_get_chain_length
        procedure, public :: push_new_state => cb_push
        procedure, public :: get_chain => cb_get_chain
        procedure, public :: reset => cb_clear_chain

        ! Private Routines
        procedure, private :: resize_buffer => cb_resize_buffer
        procedure, private :: get_buffer_length => cb_get_buffer_length
    end type

! ------------------------------------------------------------------------------
    type, extends(chain_builder) :: mcmc_sampler
        !! An implementation of the Metropolis-Hastings algorithm for the
        !! generation of a Markov chain.
        integer(int32), private :: m_accepted = 0
            !! The number of accepted steps.
    contains
        procedure, public :: sample => ms_sample 
        procedure, public :: on_acceptance => ms_on_success
        procedure, public :: on_rejection => ms_on_rejection
        procedure, public :: get_accepted_count => ms_get_num_accepted
    end type

contains
! ******************************************************************************
! MCMC_TARGET
! ------------------------------------------------------------------------------
pure function mt_get_param_count(this) result(rst)
    !! Gets the number of model parameters.
    class(mcmc_target), intent(in) :: this
        !! The mcmc_target object.
    integer(int32) :: rst
        !! The parameter count.

    rst = this%m_items%count()
end function

! ------------------------------------------------------------------------------
subroutine mt_add_param(this, x)
    !! Adds a new model parameter.
    class(mcmc_target), intent(inout) :: this
        !! The mcmc_target object.
    class(distribution), intent(in) :: x
        !! The parameter to add.

    ! Process
    call this%m_items%push(x)
end subroutine

! ------------------------------------------------------------------------------
function mt_get_param(this, i) result(rst)
    !! Gets a pointer to the stored parameter.
    class(mcmc_target), intent(in) :: this
        !! The mcmc_target object.
    integer(int32), intent(in) :: i
        !! The index of the parameter to retrieve.  If outside the bounds of the
        !! collection of parameters a null pointer is returned.
    class(distribution), pointer :: rst
        !! A pointer to the requested parameter distribution.

    ! Local Variables
    class(*), pointer :: ptr

    ! Process
    ptr => this%m_items%get(i)
    rst => null()
    select type (ptr)
    class is (distribution)
        rst => ptr
    end select
end function

! ------------------------------------------------------------------------------
function mt_likelihood(this, xdata, ydata, xc, var) result(rst)
    !! Computes the Gaussian log likelihood. Invalid variance gives NaN.
    class(mcmc_target), intent(inout) :: this
        !! The mcmc_target object.
    real(real64), intent(in), dimension(:) :: xdata
        !! An M-element array containing the independent data points.
    real(real64), intent(in), dimension(:) :: ydata
        !! An M-element array containing the dependent data points.
    real(real64), intent(in), dimension(:) :: xc
        !! An N-element array containing the model parameters.
    real(real64), intent(in) :: var
        !! Finite, strictly positive model variance; it is not silently clamped.
    real(real64) :: rst
        !! Log likelihood, or NaN for invalid/nonfinite inputs.

    rst = ieee_value(0.0d0, ieee_quiet_nan)
    if (.not.ieee_is_finite(var) .or. var <= 0.0d0) return
    rst = this%likelihood_log_variance(xdata, ydata, xc, log(var))
end function

pure elemental function log_absolute_difference(first, second) result(rst)
    !! Computes log(abs(first - second)) without overflowing the subtraction.
    real(real64), intent(in) :: first, second
        !! Finite values whose difference is required.
    real(real64) :: rst
        !! Log magnitude; equal inputs give -huge as an internal zero sentinel.
    real(real64) :: scale

    rst = -huge(1.0d0)
    if (first == second) return
    if ((first >= 0.0d0 .and. second >= 0.0d0) .or. &
        (first <= 0.0d0 .and. second <= 0.0d0)) then
        rst = log(abs(first - second))
    else
        scale = max(abs(first), abs(second))
        rst = log(scale) + log(abs(first / scale - second / scale))
    end if
end function

function mt_log_likelihood(this, xdata, ydata, xc, log_variance) result(rst)
    !! Computes the Gaussian log likelihood without exponentiating the variance.
    !! Residual magnitudes are scaled in log space before their squared reduction.
    !! Override this binding for custom likelihoods used by the sampler; overriding
    !! only likelihood changes calls to that variance-space entry point, not sampling.
    class(mcmc_target), intent(inout) :: this
        !! Target; its model workspace is resized as necessary.
    real(real64), intent(in), dimension(:) :: xdata
        !! M finite independent observations.
    real(real64), intent(in), dimension(:) :: ydata
        !! M finite dependent observations.
    real(real64), intent(in), dimension(:) :: xc
        !! Finite parameters, one per registered parameter distribution.
    real(real64), intent(in) :: log_variance
        !! Finite natural logarithm of the model variance.
    real(real64) :: rst
        !! Log likelihood; empty data give zero, invalid inputs/models give NaN,
        !! and an unrepresentably large residual penalty gives negative infinity.
    integer(int32) :: i, m
    real(real64) :: log_scale, log_penalty, normalization
    logical :: use_parallel

    m = size(xdata)
    if (size(ydata) /= m) error stop FS_ARRAY_SIZE_ERROR
    if (this%get_parameter_count() /= size(xc)) error stop FS_ARRAY_SIZE_ERROR
    rst = ieee_value(0.0d0, ieee_quiet_nan)
    if (.not.ieee_is_finite(log_variance)) return
    if (.not.all(ieee_is_finite(xdata)) .or. .not.all(ieee_is_finite(ydata))) return
    if (.not.all(ieee_is_finite(xc))) return
    if (m == 0) then
        rst = 0.0d0
        return
    end if
    if (.not.allocated(this%m_y)) then
        allocate(this%m_y(m))
    end if
    if (size(this%m_y) /= m) then
        deallocate(this%m_y)
        allocate(this%m_y(m))
    end if

    ! Evaluate the model at each data point
    call this%model(xdata, xc, this%m_y)
    if (.not.all(ieee_is_finite(this%m_y))) return

    use_parallel = m >= this%likelihood_parallel_threshold
    !$ use_parallel = use_parallel .and. .not. omp_in_parallel()
    log_scale = -huge(1.0d0)
    !$omp parallel do simd default(none) if(use_parallel) &
    !$omp shared(this, ydata, m, use_parallel) reduction(max:log_scale) schedule(static)
    do i = 1, m
        log_scale = max(log_scale, log_absolute_difference(ydata(i), this%m_y(i)))
    end do
    !$omp end parallel do simd
    rst = 0.0d0
    if (log_scale > -huge(1.0d0)) then
        !$omp parallel do simd default(none) if(use_parallel) &
        !$omp shared(this, ydata, m, log_scale, use_parallel) reduction(+:rst) schedule(static)
        do i = 1, m
            if (ydata(i) /= this%m_y(i)) rst = rst + &
                exp(2.0d0 * (log_absolute_difference(ydata(i), this%m_y(i)) - log_scale))
        end do
        !$omp end parallel do simd
        log_penalty = 2.0d0 * log_scale - log_variance + log(rst) - log(2.0d0)
        if (log_penalty > log(huge(1.0d0))) then
            rst = ieee_value(0.0d0, ieee_negative_inf)
            return
        end if
        rst = exp(log_penalty)
    end if
    normalization = 0.5d0 * (log(4.0d0 * acos(0.0d0)) + log_variance)
    if (abs(normalization) > huge(1.0d0) / real(m, real64)) then
        rst = ieee_value(0.0d0, ieee_quiet_nan)
        return
    end if
    rst = -rst - real(m, real64) * normalization
end function

pure function mt_eval_log_var_prior(this, log_variance) result(rst)
    !! Log density of log variance, including the variance-transform Jacobian.
    !! The default variance-space prior is log-normal, so its transformed density
    !! is normal with mean zero and standard deviation data_noise. Override this
    !! binding when changing the variance prior used by the sampler.
    class(mcmc_target), intent(in) :: this
        !! Target with finite, positive data_noise.
    real(real64), intent(in) :: log_variance
        !! Natural logarithm of the variance.
    real(real64) :: rst
        !! Log transformed prior; invalid prior parameters or NaN input give NaN.
    type(normal_distribution) :: prior

    prior%mean_value = 0.0d0
    prior%standard_deviation = this%data_noise
    rst = prior%log_pdf(log_variance)
end function

function mt_log_posterior(this, xdata, ydata, xc, log_variance) result(rst)
    !! Evaluates the posterior density with respect to parameters and log variance.
    !! Includes the variance-transform Jacobian through evaluate_log_variance_prior.
    !! Impossible parameter states are rejected before model evaluation.
    class(mcmc_target), intent(inout) :: this
        !! Target whose prior and log-likelihood bindings are evaluated.
    real(real64), intent(in), dimension(:) :: xdata
        !! Independent observations.
    real(real64), intent(in), dimension(:) :: ydata
        !! Dependent observations, the same size as xdata.
    real(real64), intent(in), dimension(:) :: xc
        !! Model parameters, one per registered distribution.
    real(real64), intent(in) :: log_variance
        !! Natural logarithm of the model variance.
    real(real64) :: rst
        !! Log posterior; impossible states give -infinity, invalid inputs give NaN.
    real(real64) :: prior, variance_prior

    prior = this%evaluate_prior(xc)
    variance_prior = this%evaluate_log_variance_prior(log_variance)
    rst = ieee_value(0.0d0, ieee_quiet_nan)
    if (ieee_is_nan(prior) .or. ieee_is_nan(variance_prior)) return
    rst = ieee_value(0.0d0, ieee_negative_inf)
    if (.not.ieee_is_finite(prior) .or. .not.ieee_is_finite(variance_prior)) return
    rst = this%likelihood_log_variance(xdata, ydata, xc, log_variance)
    rst = rst + prior + variance_prior
end function

! ------------------------------------------------------------------------------
pure function mt_eval_var_prior(this, x) result(rst)
    !! Computes the log prior density with respect to variance, without its
    !! log-coordinate Jacobian. Sampling uses evaluate_log_variance_prior instead.
    class(mcmc_target), intent(in) :: this
        !! The mcmc_target object.
    real(real64), intent(in) :: x
        !! The value at which to evaluate the variance prior distribution PDF.
    real(real64) :: rst
        !! Log prior density; outside support gives -infinity, invalid inputs NaN.

    ! Local Variables
    type(log_normal_distribution) :: dist

    ! Initialization
    dist%mean_value = 0.0d0
    dist%standard_deviation = this%data_noise
    
    ! Process: return log of the variance prior PDF; handle zero/negative PDF
    rst = dist%log_pdf(x)
end function

! ------------------------------------------------------------------------------
function mt_sample_var_prior(this, vc, n) result(rst)
    !! Samples the variance prior distribution for the requested number of 
    !! samples.
    class(mcmc_target), intent(inout) :: this
        !! The mcmc_target object.
    real(real64), intent(in) :: vc
        !! The prior variance term.
    integer(int32), intent(in) :: n
        !! The number of samples.
    real(real64), allocatable, dimension(:) :: rst
        !! The requested samples.

    ! Local Variables
    type(log_normal_distribution) :: dist
    real(real64) :: xmin, xmax, sigma

    ! Establish an upper bounds on the sampling region
    dist%mean_value = vc
    dist%standard_deviation = this%data_noise
    sigma = dist%standard_deviation
    if (sigma == 0.0d0) then
        sigma = 1.0d0
        dist%standard_deviation = 1.0d0
    end if
    xmax = 6.0d0 * sigma
    xmin = 1.0d-6 * xmax

    ! Process
    rst = rejection_sample(dist, n, xmin, xmax)
end function

! ------------------------------------------------------------------------------
function mt_eval_prior(this, x) result(rst)
    !! Evaluates the PDF's for each parameter and computes a probability.
    class(mcmc_target), intent(in) :: this
        !! The mcmc_target object.
    real(real64), intent(in), dimension(:) :: x
        !! An N-element array containing the values at which to evaluate each of
        !! the N parameter PDF's.
    real(real64) :: rst
        !! The resulting probability.

    ! Local Variables
    integer(int32) :: i, n
    real(real64) :: temp, p
    class(distribution), pointer :: dist
    
    ! Initialization
    n = this%get_parameter_count()

    ! Input Check
    if (size(x) /= n) error stop FS_ARRAY_SIZE_ERROR

    ! Process - use natural-log probabilities to avoid overflow/underflow issues
    temp = 0.0d0
    do i = 1, n
        ! Evaluate the distribution
        dist => this%get_parameter(i)
        if (.not.associated(dist)) error stop FS_NULL_POINTER_ERROR
        p = dist%log_pdf(x(i))
        if (.not.ieee_is_finite(p)) then
            rst = p
            return
        end if
        temp = temp + p
    end do
    rst = temp  ! return log-prior
end function

! ******************************************************************************
! MCMC_PROPOSAL
! ------------------------------------------------------------------------------
subroutine mp_gen(this, tgt, xc, xp, vc, vp)
    !! Creates a new sample proposal.
    class(mcmc_proposal), intent(inout) :: this
        !! The mcmc_proposal object.
    class(mcmc_target), intent(inout) :: tgt
        !! The mcmc_target object.
    real(real64), intent(in), dimension(:) :: xc
        !! An N-element array containing the existing parameter estimates.
    real(real64), intent(out), dimension(:) :: xp
        !! An N-element array where the proposed parameters will be output.
    real(real64), intent(in) :: vc
        !! The current variance (noise) term value.
    real(real64), intent(out) :: vp
        !! The proposed variance (noise) value.

    ! Parameters
    integer(int32), parameter :: nsamples = 1

    ! Local Variables
    integer(int32) :: i, n
    real(real64) :: samples(nsamples), sigma, mu, xmax, xmin, mx, mn, rng(2)
    real(real64) :: u1, u2, z, u1v, u2v, zv, pscale
    class(distribution), pointer :: dist
    
    ! Initialization
    n = tgt%get_parameter_count()

    ! Input Checking
    if (size(xp) /= n) error stop FS_ARRAY_SIZE_ERROR

    ! Use a symmetric Gaussian random-walk proposal for each parameter.
    ! Proposals are: xp(i) = xc(i) + scale * sigma_i * Normal(0,1)
    ! For the variance term we propose on log-scale: vp = vc + scale * Normal(0,1)
    do i = 1, n
        ! Get the parameter distribution for scale information
        dist => tgt%get_parameter(i)
        if (.not.associated(dist)) error stop FS_NULL_POINTER_ERROR

        ! Recenter is not needed for RW proposals; ignore recenter flag
        sigma = dist%variance()
        if (.not.ieee_is_finite(sigma) .or. sigma <= 0.0d0) error stop FS_INVALID_INPUT_ERROR
        sigma = sqrt(sigma)

        ! Generate a normal(0,1) using Box-Muller
        do
            call random_number(u1)
            if (u1 > 0.0d0) exit
        end do
        call random_number(u2)
        z = sqrt(-2.0d0 * log(u1)) * cos(2.0d0 * acos(0.0d0) * u2)

        if (allocated(this%m_param_scale)) then
            if (i <= size(this%m_param_scale)) then
                pscale = this%m_param_scale(i)
            else
                pscale = 1.0d0
            end if
        else
            pscale = 1.0d0
        end if
        xp(i) = xc(i) + this%m_scale * pscale * sigma * z
    end do

    ! Propose log-variance (vc is the current log-variance)
    ! generate another normal variate for variance step
    do
        call random_number(u1v)
        if (u1v > 0.0d0) exit
    end do
    call random_number(u2v)
    zv = sqrt(-2.0d0 * log(u1v)) * cos(2.0d0 * acos(0.0d0) * u2v)
    vp = vc + this%m_scale * zv
end subroutine

! ------------------------------------------------------------------------------
pure function mp_get_recenter(this) result(rst)
    !! Gets a value determining if the parameter distributions should be
    !! recentered about the last stored position upon sampling.
    class(mcmc_proposal), intent(in) :: this
        !! The mcmc_proposal object.
    logical :: rst
        !! True if recentering is to be allowed; else, false.

    rst = this%m_recenter
end function

! ------------------------------------------------------------------------------
subroutine mp_set_recenter(this, x)
    !! Sets a value determining if the parameter distributions should be 
    !! recentered about the last stored position upon sampling.
    class(mcmc_proposal), intent(inout) :: this
        !! The mcmc_proposal object.
    logical, intent(in) :: x
        !! True if recentering is to be allowed; else, false.

    this%m_recenter = x
end subroutine

! ------------------------------------------------------------------------------
pure function mp_get_scale(this) result(rst)
    !! Gets the proposal scale multiplier.
    class(mcmc_proposal), intent(in) :: this
    real(real64) :: rst
    rst = this%m_scale
end function

! ------------------------------------------------------------------------------
subroutine mp_set_scale(this, x)
    !! Sets the proposal scale multiplier.
    class(mcmc_proposal), intent(inout) :: this
    real(real64), intent(in) :: x
    if (.not.ieee_is_finite(x) .or. x <= 0.0d0) error stop FS_INVALID_INPUT_ERROR
    this%m_scale = x
end subroutine

! ------------------------------------------------------------------------------
pure function mp_get_param_scale(this, idx) result(rst)
    class(mcmc_proposal), intent(in) :: this
    integer(int32), intent(in) :: idx
    real(real64) :: rst
    if (allocated(this%m_param_scale)) then
        if (idx >= 1 .and. idx <= size(this%m_param_scale)) then
            rst = this%m_param_scale(idx)
        else
            rst = this%m_scale
        end if
    else
        rst = this%m_scale
    end if
end function

! ------------------------------------------------------------------------------
subroutine mp_set_param_scale(this, idx, val)
    class(mcmc_proposal), intent(inout) :: this
    integer(int32), intent(in) :: idx
    real(real64), intent(in) :: val
    integer(int32) :: n
    if (idx < 1) error stop FS_INVALID_INPUT_ERROR
    if (.not.ieee_is_finite(val) .or. val <= 0.0d0) error stop FS_INVALID_INPUT_ERROR
    if (.not.allocated(this%m_param_scale)) then
        n = max(1_int32, idx)
        allocate(this%m_param_scale(n))
        this%m_param_scale = 1.0d0
    else
        if (idx > size(this%m_param_scale)) then
            n = idx
            this%m_param_scale = [this%m_param_scale, spread(1.0d0, 1, n - size(this%m_param_scale))]
        end if
    end if
    this%m_param_scale(idx) = val
end subroutine

! ------------------------------------------------------------------------------
subroutine mp_set_param_scales(this, arr)
    class(mcmc_proposal), intent(inout) :: this
    real(real64), intent(in), dimension(:) :: arr
    if (.not.all(ieee_is_finite(arr)) .or. any(arr <= 0.0d0)) error stop FS_INVALID_INPUT_ERROR
    if (.not.allocated(this%m_param_scale)) then
        allocate(this%m_param_scale(size(arr)))
    else
        if (size(this%m_param_scale) /= size(arr)) then
            deallocate(this%m_param_scale)
            allocate(this%m_param_scale(size(arr)))
        end if
    end if
    this%m_param_scale = arr
end subroutine


! ******************************************************************************
! CHAIN_BUILDER
! ------------------------------------------------------------------------------
pure function cb_get_nvars(this) result(rst)
    !! Gets the number of state variables.
    class(chain_builder), intent(in) :: this
        !! The chain_builder object.
    integer(int32) :: rst
        !! The number of state variables.

    rst = this%m_numVars
end function

! ------------------------------------------------------------------------------
pure function cb_get_chain_length(this) result(rst)
    !! Gets the length of the chain (number of stored state variables).
    class(chain_builder), intent(in) :: this
        !! The chain_builder object.
    integer(int32) :: rst
        !! The chain length.

    rst = this%m_bufferSize
end function

! ------------------------------------------------------------------------------
subroutine cb_resize_buffer(this)
    !! Resizes the buffer to accept more states.
    class(chain_builder), intent(inout) :: this
        !! The chain_builder object.

    ! Local Variables
    integer(int32) :: m, n, mOld
    real(real64), allocatable, dimension(:,:) :: copy
    
    ! Initialization
    m = this%initial_iteration_estimate
    n = this%get_state_variable_count()

    ! Is this the first time?
    if (.not.allocated(this%m_buffer)) then
        allocate(this%m_buffer(m, n))
    end if

    ! If we're here, then we need to create a copy and go from there
    m = size(this%m_buffer, 1)
    mOld = m
    allocate(copy(m, n), source = this%m_buffer)
    deallocate(this%m_buffer)
    m = m + this%initial_iteration_estimate
    allocate(this%m_buffer(m, n))
    this%m_buffer(1:mOld,:) = copy
    deallocate(copy)
end subroutine

! ------------------------------------------------------------------------------
pure function cb_get_buffer_length(this) result(rst)
    !! Gets the actual length of the buffer.  This value will likely exceed the
    !! actual number of items in the chain.
    class(chain_builder), intent(in) :: this
        !! The chain_builder object.
    integer(int32) :: rst
        !! The actual buffer length.

    if (allocated(this%m_buffer)) then
        rst = size(this%m_buffer, 1)
    else
        rst = 0
    end if
end function

! ------------------------------------------------------------------------------
subroutine cb_push(this, x)
    !! Pushes a new set of state variables onto the buffer.
    class(chain_builder), intent(inout) :: this
        !! The chain_builder object.
    real(real64), intent(in), dimension(:) :: x
        !! The new N-element state array.

    ! Local Variables
    integer(int32) :: n, n1, nbuffer, nvars
    
    ! Initialization
    n = this%get_chain_length()
    n1 = n + 1
    nbuffer = this%get_buffer_length()
    nvars = size(x)

    ! If this is the first time, ensure the collection is initialized
    if (this%get_state_variable_count() == 0) then
        this%m_numVars = nvars
    end if

    ! Input Checking
    if (nvars /= this%get_state_variable_count()) error stop FS_ARRAY_SIZE_ERROR

    ! Resize the buffer, if necessary
    if (n == 0 .or. n == nbuffer) then
        call this%resize_buffer()
    end if

    ! Store the new state
    this%m_buffer(n1,:) = x
    this%m_bufferSize = this%m_bufferSize + 1
end subroutine

! ------------------------------------------------------------------------------
pure function cb_get_chain(this, bin) result(rst)
    !! Gets a copy of the stored Markov chain.
    class(chain_builder), intent(in) :: this
        !! The chain_builder object.
    real(real64), intent(in), optional :: bin
        !! An optional input allowing for a burn-in region.  The parameter
        !! represents the amount (percentage-based) of the overall chain to 
        !! disregard as "burn-in" values.  The value shoud exist on [0, 1).
        !! The default value is 0 such that no values are disregarded.
    real(real64), allocatable, dimension(:,:) :: rst
        !! The resulting chain with each parameter represented by a column.

    ! Local Variables
    integer(int32) :: npts, nvar, nstart, n
    
    ! Initialization
    npts = this%get_chain_length()
    n = npts
    nvar = this%get_state_variable_count()
    if (present(bin)) then
        nstart = int(floor(bin * real(n))) + 1
        if (nstart < 1) nstart = 1
        if (nstart > n) then
            ! No samples after burn-in
            allocate(rst(0, nvar))
        end if
        npts = n - nstart + 1
    else
        nstart = 1
    end if

    ! Process
    allocate(rst(npts, nvar), source = this%m_buffer(nstart:n,1:nvar))
end function

! ------------------------------------------------------------------------------
subroutine cb_clear_chain(this)
    !! Resets the object and clears out the buffer storing the chain values.
    class(chain_builder), intent(inout) :: this
        !! The chain_builder object.

    ! Clear the buffer
    this%m_bufferSize = 0
    this%m_numVars = 0
end subroutine

! ******************************************************************************
! MCMC_SAMPLER
! ------------------------------------------------------------------------------
subroutine sample_chains(samplers, xdata, ydata, proposals, targets, niter)
    !! Samples independent chains concurrently using OpenMP.
    !! Initialize each target independently, with no shared mutable state in
    !! distributions, models, proposals, or sampler callbacks.  The compiler
    !! runtime must provide thread-safe RANDOM_NUMBER (as GNU Fortran does).
    !! Random streams and reproducibility depend on the compiler runtime and
    !! thread scheduling; no per-chain seed guarantee is provided.
    class(mcmc_sampler), intent(inout) :: samplers(:)
        !! One sampler per chain; existing samples are retained as in sample.
    real(real64), intent(in) :: xdata(:), ydata(:)
        !! Shared, read-only observations.
    class(mcmc_proposal), intent(inout) :: proposals(:)
        !! One independently configured proposal per chain.
    class(mcmc_target), intent(inout) :: targets(:)
        !! One independently initialized target per chain.
    integer(int32), intent(in), optional :: niter
        !! Iterations per chain; defaults to 10,000.

    integer(int32) :: chain_index, nchains
    logical :: use_parallel

    nchains = size(samplers)
    if (size(proposals) /= nchains .or. size(targets) /= nchains) &
        error stop FS_ARRAY_SIZE_ERROR
    if (size(xdata) /= size(ydata)) error stop FS_ARRAY_SIZE_ERROR

    use_parallel = nchains > 1
    !$ use_parallel = use_parallel .and. .not. omp_in_parallel()
    !$omp parallel do default(none) if(use_parallel) schedule(static) &
    !$omp shared(samplers, xdata, ydata, proposals, targets, niter, nchains, use_parallel)
    do chain_index = 1, nchains
        call samplers(chain_index)%sample(xdata, ydata, proposals(chain_index), &
            targets(chain_index), niter)
    end do
    !$omp end parallel do
end subroutine

! ------------------------------------------------------------------------------
subroutine ms_on_success(this, iter, alpha, xc, xp)
    !! Currently, this routine does nothing and is a placeholder for the user
    !! that inherits this class to provide functionallity upon acceptance of
    !! a proposed value.
    class(mcmc_sampler), intent(inout) :: this
        !! The mcmc_sampler object.
    integer(int32), intent(in) :: iter
        !! The current iteration number.
    real(real64), intent(in) :: alpha
        !! The proposal probabilty term used for acceptance criteria.
    real(real64), intent(in), dimension(:) :: xc
        !! An N-element array containing the current state variables.
    real(real64), intent(in), dimension(size(xc)) :: xp
        !! An N-element array containing the proposed state variables that
        !! were just accepted.
end subroutine

! ------------------------------------------------------------------------------
subroutine ms_on_rejection(this, iter, alpha, xc, xp)
    !! Currently, this routine does nothing and is a placeholder for the user
    !! that inherits this class to provide functionallity upon rejection of
    !! a proposed value.
    class(mcmc_sampler), intent(inout) :: this
        !! The mcmc_sampler object.
    integer(int32), intent(in) :: iter
        !! The current iteration number.
    real(real64), intent(in) :: alpha
        !! The proposal probabilty term used for acceptance criteria.
    real(real64), intent(in), dimension(:) :: xc
        !! An N-element array containing the current state variables.
    real(real64), intent(in), dimension(size(xc)) :: xp
        !! An N-element array containing the proposed state variables that
        !! were just rejected.
end subroutine

! ------------------------------------------------------------------------------
pure function ms_get_num_accepted(this) result(rst)
    !! Gets the number of accepted steps.
    class(mcmc_sampler), intent(in) :: this
        !! The mcmc_sampler object.
    integer(int32) :: rst
        !! The number of accepted steps.
    rst = this%m_accepted
end function

! ------------------------------------------------------------------------------
subroutine ms_sample(this, xdata, ydata, prop, tgt, niter)
    !! Samples the distribution using the Metropolis-Hastings approach.
    class(mcmc_sampler), intent(inout) :: this
        !! The mcmc_sampler object.
    real(real64), intent(in), dimension(:) :: xdata
        !! An M-element array containing the independent coordinate data points.
    real(real64), intent(in), dimension(:) :: ydata
        !! An M-element array containing the dependent coordinate data points.
    class(mcmc_proposal), intent(inout) :: prop
        !! A proposal generation object.
    class(mcmc_target), intent(inout) :: tgt
        !! An mcmc_target-based object containing the model and allowing for
        !! evaluation of likelihoods.
    integer(int32), intent(in), optional :: niter
        !! An optional input defining the number of iterations to take.  The
        !! default is 10,000.

    ! Local Variables
    integer(int32) :: i, n, n1, npts, m
    integer(int32) :: last_accept_count, adapt_interval
    real(real64) :: alpha
    real(real64) :: pc_log, pp_log, delta, u
    real(real64) :: adapt_gain, target_accept, accept_rate, cur_scale, new_scale
    real(real64), allocatable, dimension(:) :: buffer, xc
    class(distribution), pointer :: dist
    
    ! Initialization
    if (present(niter)) then
        npts = niter
    else
        npts = this%initial_iteration_estimate
    end if
    m = size(xdata)
    n = tgt%get_parameter_count()
    n1 = n + 1  ! include the variance term (we store log-variance in the last element)
    this%m_accepted = 0

    ! Input Checking
    if (size(ydata) /= m) error stop FS_ARRAY_SIZE_ERROR
    if (npts < 1) error stop FS_INVALID_INPUT_ERROR
    if (.not.all(ieee_is_finite(xdata)) .or. .not.all(ieee_is_finite(ydata))) &
        error stop FS_INVALID_INPUT_ERROR

    ! Memory Allocations
    allocate(buffer(n1), xc(n1), source = 0.0d0)

    ! Get an initial starting point based upon the prior means
    do i = 1, n
        dist => tgt%get_parameter(i)
        if (.not.associated(dist)) error stop FS_NULL_POINTER_ERROR
        xc(i) = dist%mean()
    end do
    ! Initialize log-variance to 0 (i.e., variance = 1.0)
    xc(n1) = 0.0d0

    ! Compute initial log-posterior for the starting state
    pc_log = tgt%log_posterior(xdata, ydata, xc(1:n), xc(n1))
    if (.not.ieee_is_finite(pc_log)) error stop FS_INVALID_INPUT_ERROR

    ! Store the starting location
    call this%push_new_state(xc)

    ! Setup adaptation parameters
    adapt_interval = 100
    adapt_gain = 0.05d0
    target_accept = 0.234d0
    last_accept_count = 0

    ! Process
    do i = 2, npts
        ! Create a proposal (note: xc(1:n) are parameters, xc(n1) is log-variance)
        call prop%generate_sample(tgt, xc(1:n), buffer(1:n), xc(n1), buffer(n1))

        pp_log = tgt%log_posterior(xdata, ydata, buffer(1:n), buffer(n1))

        ! MH acceptance in log-domain
        delta = ieee_value(0.0d0, ieee_negative_inf)
        if (ieee_is_finite(pp_log)) then
            if (pp_log >= pc_log) then
                delta = 0.0d0
            else if (pc_log > 0.0d0 .and. pp_log < 0.0d0) then
                if (pp_log >= -huge(1.0d0) + pc_log) delta = pp_log - pc_log
            else
                delta = pp_log - pc_log
            end if
        end if
        alpha = exp(min(0.0d0, delta))
        do
            call random_number(u)
            if (u > 0.0d0) exit
        end do
        if (log(u) <= delta) then
            ! Accept
            call this%push_new_state(buffer)

            this%m_accepted = this%m_accepted + 1

            ! Call user hook
            call this%on_acceptance(i, alpha, xc, buffer)
            xc = buffer
            pc_log = pp_log
        else
            ! Reject: keep current state
            call this%push_new_state(xc)

            call this%on_rejection(i, alpha, xc, buffer)
        end if

        ! Adapt proposal scale every adapt_interval iterations (simple global adaptation)
        if (mod(i, adapt_interval) == 0) then
            accept_rate = real(this%m_accepted - last_accept_count, real64) / real(adapt_interval, real64)
            last_accept_count = this%m_accepted
            ! Update global scale multiplicatively toward target acceptance
            cur_scale = prop%get_scale()
            new_scale = cur_scale * exp(adapt_gain * (accept_rate - target_accept))
            call prop%set_scale(new_scale)
        end if
    end do
end subroutine

! ------------------------------------------------------------------------------
end module